# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import numpy as np

from experimental.prophet.graph import graph
from experimental.prophet.graph.shape import Shape
from experimental.prophet.graph.util import to_data_type_proto


def shared(dtype, value=None, shape=None, node_name=None, namespace=None):
  shape = Shape(shape) if shape is not None else Shape.Any()
  attrs = {'dtype': to_data_type_proto(dtype), 'shape': shape.tuple}
  if value is not None:
    value = np.asarray(value, dtype=dtype)
    attrs['value'] = value.tolist()

  node = graph.create_node('SharedVariable', [], name=node_name, namespace=namespace, attrs=attrs)
  return node.outputs[0]


def assign(shared_var, input_var, node_name=None, namespace=None):
  assert shared_var.node.op in ('SharedVariable', 'SharedVariableAssign')
  shvar_node = shared_var.node
  attrs = {
      'dtype': shvar_node.attrs['dtype'],
      'shape': shvar_node.attrs['shape'],
      'shvar': shvar_node.attrs.get('shvar', shvar_node.name)
  }

  input_var = graph.as_variable(input_var)
  assert to_data_type_proto(input_var.dtype) == attrs['dtype']

  node = graph.create_node('SharedVariableAssign', [input_var, shvar_node.outputs[1]],
                           name=node_name,
                           namespace=namespace,
                           attrs=attrs)
  return node.outputs[0]
